import logging
import sys
import os
from datetime import datetime, timedelta
from sqlalchemy import create_engine, Column, Integer, String, Float, Text, JSON, DateTime
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, Session

# 直接定义数据库URL和创建engine
DATABASE_URL = "sqlite:///./pilot_manage.db"
engine = create_engine(DATABASE_URL, connect_args={"check_same_thread": False})
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
Base = declarative_base()

# 直接在脚本中定义NoFlyZone模型，确保使用正确的Base和字段
class NoFlyZone(Base):
    __tablename__ = "no_fly_zones"
    
    id = Column(Integer, primary_key=True, index=True)
    name = Column(String(255), nullable=False)
    description = Column(Text)
    zone_type = Column(String(50), nullable=False)  # circle 或 polygon
    level = Column(String(50))  # 禁飞区级别
    color = Column(String(20))  # 显示颜色
    lat = Column(Float)  # 圆心纬度
    lng = Column(Float)  # 圆心经度
    radius = Column(Float)  # 半径(米)
    polygon_points = Column(JSON)  # 多边形点坐标
    height = Column(String(50))  # 高度描述
    min_altitude = Column(Float, default=0)  # 最低高度
    max_altitude = Column(Float, default=1000)  # 最高高度
    sub_areas = Column(JSON)  # 子区域列表
    start_time = Column(DateTime)  # 开始时间
    end_time = Column(DateTime)  # 结束时间
    created_at = Column(DateTime, default=datetime.now)
    updated_at = Column(DateTime, default=datetime.now, onupdate=datetime.now)

# 获取数据库会话的函数
def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 预设的禁飞区数据 - 适配DJI API格式
PRESET_NO_FLY_ZONES = [
    # 圆形禁飞区示例 - 机场禁飞区
    {
        "name": "机场禁飞区",
        "description": "机场周边5公里禁飞区",
        "zone_type": "circle",
        "lat": 39.5000,  # 示例坐标 - 圆心纬度
        "lng": 116.4000, # 示例坐标 - 圆心经度
        "radius": 5000,  # 5公里
        "min_altitude": 0,
        "max_altitude": 1000,
        "level": "airport",
        "color": "#FF0000",
        "height": "0m-1000m"
    },
    # 多边形禁飞区示例 - 军事基地禁飞区
    {
        "name": "军事基地禁飞区",
        "description": "军事基地多边形禁飞区",
        "zone_type": "polygon",
        "polygon_points": [
            [116.3000, 39.9000],  # [lng, lat] 格式
            [116.3500, 39.9000],
            [116.3500, 39.8500],
            [116.3000, 39.8500]
        ],
        "min_altitude": 0,
        "max_altitude": 1000,
        "level": "military",
        "color": "#FF3333",
        "height": "0m-1000m"
    },
    # 临时禁飞区示例 - 重大活动临时禁飞区
    {
        "name": "重大活动临时禁飞区",
        "description": "重大活动期间的临时禁飞区",
        "zone_type": "circle",
        "lat": 39.7000,  # 示例坐标
        "lng": 116.5000, # 示例坐标
        "radius": 2000,  # 2公里
        "min_altitude": 0,
        "max_altitude": 500,
        "level": "temporary",
        "color": "#FF9900",
        "height": "0m-500m",
        "start_time": datetime.now(),
        "end_time": datetime.now() + timedelta(days=7)  # 有效期7天
    },
    # 多边形禁飞区示例2 - 核电站禁飞区
    {
        "name": "核电站禁飞区",
        "description": "核电站周边多边形禁飞区",
        "zone_type": "polygon",
        "polygon_points": [
            [116.7000, 39.6000],
            [116.7500, 39.6000],
            [116.7500, 39.5500],
            [116.7000, 39.5500]
        ],
        "min_altitude": 0,
        "max_altitude": 1000,
        "level": "nuclear",
        "color": "#CC0000",
        "height": "0m-1000m"
    },
    # 政府建筑禁飞区 - 包含子区域示例
    {
        "name": "政府建筑禁飞区",
        "description": "政府重要建筑周边禁飞区",
        "zone_type": "circle",
        "lat": 39.8000,  # 示例坐标
        "lng": 116.6000, # 示例坐标
        "radius": 1000,  # 1公里
        "min_altitude": 0,
        "max_altitude": 300,
        "level": "government",
        "color": "#FF6600",
        "height": "0m-300m",
        # 添加子区域示例
        "sub_areas": [
            {
                "lat": 39.8050,
                "lng": 116.6050,
                "radius": 500,
                "color": "#FF0000",
                "level": "government-core",
                "height": "0m-200m"
            }
        ]
    }
]

def init_no_fly_zones():
    """初始化禁飞区数据"""
    logger.info("开始初始化禁飞区数据...")
    
    # 先删除现有表，确保使用新的模型定义
    NoFlyZone.__table__.drop(bind=engine, checkfirst=True)
    logger.info("已删除现有禁飞区表")
    
    # 创建数据库表
    Base.metadata.create_all(bind=engine)
    logger.info("数据库表创建成功")
    
    # 获取数据库会话
    db: Session = next(get_db())
    
    try:
        # 统计已存在的禁飞区数量
        existing_count = db.query(NoFlyZone).count()
        logger.info(f"当前数据库中已有禁飞区数量: {existing_count}")
        
        # 添加预设的禁飞区数据
        added_count = 0
        for zone_data in PRESET_NO_FLY_ZONES:
            # 检查是否已存在同名禁飞区
            existing_zone = db.query(NoFlyZone).filter(NoFlyZone.name == zone_data["name"]).first()
            
            if existing_zone:
                logger.info(f"禁飞区 '{zone_data['name']}' 已存在，跳过...")
            else:
                # 创建新禁飞区
                no_fly_zone = NoFlyZone(**zone_data)
                db.add(no_fly_zone)
                added_count += 1
                logger.info(f"添加禁飞区: {zone_data['name']}")
        
        # 提交事务
        db.commit()
        logger.info(f"成功添加 {added_count} 个禁飞区数据")
        
    except Exception as e:
        logger.error(f"初始化禁飞区数据时出错: {str(e)}")
        db.rollback()
        raise
    finally:
        # 关闭数据库会话
        db.close()

if __name__ == "__main__":
    init_no_fly_zones()